"""Cairo 代码分析节点"""

import re
import ast
from typing import Dict, List, Any, Optional, Tuple
from core.graph.chat_graph import ChatState, CodeAnalysisResult


class CairoCodeAnalyzer:
    """Cairo 代码分析器"""
    
    def __init__(self):
        self.cairo_keywords = {
            'fn', 'struct', 'trait', 'impl', 'mod', 'contract', 'enum',
            'let', 'mut', 'const', 'use', 'pub', 'if', 'else', 'loop',
            'while', 'for', 'match', 'return', 'assert', 'panic'
        }
        
        self.cairo_types = {
            'felt252', 'u8', 'u16', 'u32', 'u64', 'u128', 'u256',
            'bool', 'Array', 'Span', 'ByteArray', 'ContractAddress'
        }
        
        self.starknet_imports = {
            'starknet', 'get_caller_address', 'get_contract_address',
            'storage_read', 'storage_write', 'emit_event'
        }

    def extract_code_blocks(self, text: str) -> List[str]:
        """提取代码块"""
        # 提取 ``` 包围的代码块
        code_blocks = re.findall(r'```(?:cairo|rust)?\n?(.*?)```', text, re.DOTALL)
        
        # 如果没有代码块，检查是否有内联代码
        if not code_blocks:
            # 检查是否包含 Cairo 关键字，可能是内联代码
            if any(keyword in text for keyword in self.cairo_keywords):
                code_blocks = [text]
        
        return [block.strip() for block in code_blocks if block.strip()]

    def analyze_syntax(self, code: str) -> List[Dict[str, Any]]:
        """分析语法错误"""
        errors = []
        
        # 检查括号匹配
        brackets = {'(': ')', '[': ']', '{': '}'}
        stack = []
        
        for i, char in enumerate(code):
            if char in brackets:
                stack.append((char, i))
            elif char in brackets.values():
                if not stack:
                    errors.append({
                        'type': 'unmatched_bracket',
                        'message': f'Unmatched closing bracket "{char}" at position {i}',
                        'position': i,
                        'severity': 'error'
                    })
                else:
                    open_bracket, _ = stack.pop()
                    if brackets[open_bracket] != char:
                        errors.append({
                            'type': 'mismatched_bracket',
                            'message': f'Mismatched bracket: expected "{brackets[open_bracket]}", got "{char}"',
                            'position': i,
                            'severity': 'error'
                        })
        
        # 检查未闭合的括号
        for bracket, pos in stack:
            errors.append({
                'type': 'unclosed_bracket',
                'message': f'Unclosed bracket "{bracket}" at position {pos}',
                'position': pos,
                'severity': 'error'
            })
        
        # 检查基本语法模式
        lines = code.split('\n')
        for line_num, line in enumerate(lines, 1):
            line = line.strip()
            if not line or line.startswith('//'):
                continue
                
            # 检查函数定义
            if line.startswith('fn '):
                if not re.match(r'fn\s+\w+\s*\([^)]*\)\s*(->\s*\w+)?\s*\{?', line):
                    errors.append({
                        'type': 'invalid_function_syntax',
                        'message': f'Invalid function syntax at line {line_num}',
                        'line': line_num,
                        'severity': 'warning'
                    })
            
            # 检查结构体定义
            if line.startswith('struct '):
                if not re.match(r'struct\s+\w+\s*\{?', line):
                    errors.append({
                        'type': 'invalid_struct_syntax',
                        'message': f'Invalid struct syntax at line {line_num}',
                        'line': line_num,
                        'severity': 'warning'
                    })
        
        return errors

    def analyze_ast_structure(self, code: str) -> Optional[Dict[str, Any]]:
        """分析 AST 结构（简化版）"""
        try:
            ast_info = {
                'functions': [],
                'structs': [],
                'traits': [],
                'impls': [],
                'imports': [],
                'contracts': []
            }
            
            lines = code.split('\n')
            current_block = None
            brace_count = 0
            
            for line_num, line in enumerate(lines, 1):
                line = line.strip()
                if not line or line.startswith('//'):
                    continue
                
                # 计算大括号层级
                brace_count += line.count('{') - line.count('}')
                
                # 函数定义
                fn_match = re.match(r'fn\s+(\w+)\s*\(([^)]*)\)\s*(->\s*([^{]+))?\s*\{?', line)
                if fn_match:
                    func_name = fn_match.group(1)
                    params = fn_match.group(2).strip() if fn_match.group(2) else ""
                    return_type = fn_match.group(4).strip() if fn_match.group(4) else "void"
                    
                    ast_info['functions'].append({
                        'name': func_name,
                        'parameters': params,
                        'return_type': return_type,
                        'line': line_num,
                        'visibility': 'pub' if 'pub fn' in line else 'private'
                    })
                
                # 结构体定义
                struct_match = re.match(r'struct\s+(\w+)\s*\{?', line)
                if struct_match:
                    struct_name = struct_match.group(1)
                    ast_info['structs'].append({
                        'name': struct_name,
                        'line': line_num,
                        'visibility': 'pub' if 'pub struct' in line else 'private'
                    })
                
                # Trait 定义
                trait_match = re.match(r'trait\s+(\w+)', line)
                if trait_match:
                    trait_name = trait_match.group(1)
                    ast_info['traits'].append({
                        'name': trait_name,
                        'line': line_num
                    })
                
                # Impl 定义
                impl_match = re.match(r'impl\s+(\w+)', line)
                if impl_match:
                    impl_name = impl_match.group(1)
                    ast_info['impls'].append({
                        'name': impl_name,
                        'line': line_num
                    })
                
                # 导入语句
                if line.startswith('use '):
                    import_path = line[4:].rstrip(';').strip()
                    ast_info['imports'].append({
                        'path': import_path,
                        'line': line_num
                    })
                
                # 合约定义
                if '#[contract]' in line or '@contract' in line:
                    ast_info['contracts'].append({
                        'line': line_num,
                        'type': 'contract_annotation'
                    })
            
            return ast_info
            
        except Exception as e:
            return {'error': str(e)}

    def calculate_complexity(self, code: str) -> float:
        """计算代码复杂度"""
        complexity = 1  # 基础复杂度
        
        # 控制流语句增加复杂度
        control_flow = ['if', 'else', 'while', 'for', 'loop', 'match']
        for keyword in control_flow:
            complexity += code.count(keyword)
        
        # 函数调用增加复杂度
        complexity += len(re.findall(r'\w+\s*\(', code))
        
        # 嵌套层级
        max_nesting = 0
        current_nesting = 0
        for char in code:
            if char == '{':
                current_nesting += 1
                max_nesting = max(max_nesting, current_nesting)
            elif char == '}':
                current_nesting -= 1
        
        complexity += max_nesting * 2
        
        # 归一化到 0-10 范围
        return min(complexity / 10.0, 10.0)

    def detect_dependencies(self, code: str) -> List[str]:
        """检测依赖项"""
        dependencies = []
        
        # 从 use 语句中提取
        use_statements = re.findall(r'use\s+([^;]+);', code)
        for use_stmt in use_statements:
            # 提取主要模块名
            main_module = use_stmt.split('::')[0].strip()
            if main_module not in dependencies:
                dependencies.append(main_module)
        
        # 检测 Starknet 相关依赖
        if any(keyword in code for keyword in self.starknet_imports):
            if 'starknet' not in dependencies:
                dependencies.append('starknet')
        
        return dependencies

    def detect_security_issues(self, code: str) -> List[Dict[str, Any]]:
        """检测安全问题"""
        issues = []
        
        # 1. 检测整数下溢/溢出漏洞
        if re.search(r'(\w+)\s*-\s*(\w+)', code):
            # 检查是否有减法操作但没有检查
            if not re.search(r'assert|require|check|>=|<=', code):
                issues.append({
                    'type': 'integer_underflow',
                    'message': '🚨 严重安全漏洞：整数下溢风险',
                    'severity': 'critical',
                    'description': '代码中存在减法操作但没有检查余额是否足够，可能导致整数下溢',
                    'suggestion': '添加余额检查：assert(balance >= amount, "Insufficient balance")',
                    'cwe': 'CWE-191',
                    'impact': '攻击者可能通过下溢获得巨额余额'
                })
        
        # 2. 检测重入攻击漏洞
        if re.search(r'call_contract|invoke', code) and re.search(r'storage_write|\.write\(', code):
            # 检查是否在外部调用后修改状态
            lines = code.split('\n')
            call_line = -1
            write_line = -1
            for i, line in enumerate(lines):
                if re.search(r'call_contract|invoke', line):
                    call_line = i
                if re.search(r'storage_write|\.write\(', line):
                    write_line = i
            
            if call_line >= 0 and write_line >= 0 and write_line > call_line:
                issues.append({
                    'type': 'reentrancy',
                    'message': '🚨 重入攻击漏洞',
                    'severity': 'critical',
                    'description': '在外部调用后修改状态，可能导致重入攻击',
                    'suggestion': '使用检查-效果-交互模式：先修改状态，再进行外部调用',
                    'cwe': 'CWE-841'
                })
        
        # 3. 检测访问控制缺失
        if re.search(r'#\[external\(v0\)\]|#\[external\]', code):
            # 这是一个外部函数
            if not re.search(r'get_caller_address|only_owner|require|assert.*caller|msg\.sender', code):
                issues.append({
                    'type': 'missing_access_control',
                    'message': '🔓 访问控制缺失',
                    'severity': 'high',
                    'description': '外部函数缺少访问控制，任何人都可以调用',
                    'suggestion': '添加访问控制：let caller = get_caller_address(); assert(caller == owner, "Unauthorized")',
                    'cwe': 'CWE-284'
                })
        
        # 4. 检测余额检查缺失
        if re.search(r'transfer|send|withdraw', code, re.IGNORECASE):
            if not re.search(r'assert.*>=|require.*>=|check.*balance|\.read\(.*\)\s*>=', code):
                issues.append({
                    'type': 'insufficient_balance_check',
                    'message': '💰 余额检查不足',
                    'severity': 'high',
                    'description': '转账函数缺少充分的余额检查',
                    'suggestion': '添加余额检查：assert(self.balances.read(caller) >= amount, "Insufficient balance")',
                    'cwe': 'CWE-682'
                })
        
        # 5. 检测状态一致性问题
        if re.search(r'\.write\(.*,.*-', code) and not re.search(r'\.write\(.*,.*\+', code):
            issues.append({
                'type': 'state_inconsistency',
                'message': '⚖️ 状态不一致风险',
                'severity': 'medium',
                'description': '只减少发送方余额，没有增加接收方余额，可能导致代币丢失',
                'suggestion': '确保状态变更的原子性：减少发送方余额的同时增加接收方余额',
                'cwe': 'CWE-362'
            })
        
        # 6. 检测事件缺失
        if re.search(r'transfer|mint|burn', code, re.IGNORECASE):
            if not re.search(r'emit|event', code, re.IGNORECASE):
                issues.append({
                    'type': 'missing_events',
                    'message': '📢 缺少事件记录',
                    'severity': 'medium',
                    'description': '重要操作缺少事件记录，影响透明度和可追踪性',
                    'suggestion': '添加事件：self.emit(Transfer { from: caller, to, amount })',
                    'cwe': 'CWE-778'
                })
        
        # 7. 检测零地址检查
        if re.search(r'to:|to\s*:', code):
            if not re.search(r'assert.*!=.*0|require.*!=.*0|check.*zero', code):
                issues.append({
                    'type': 'zero_address_check',
                    'message': '🚫 零地址检查缺失',
                    'severity': 'medium',
                    'description': '没有检查接收地址是否为零地址',
                    'suggestion': '添加零地址检查：assert(to != contract_address_const::<0>(), "Cannot transfer to zero address")',
                    'cwe': 'CWE-20'
                })
        
        # 8. 检测整数溢出
        if re.search(r'\+.*u256|u256.*\+', code):
            if not re.search(r'checked_add|overflow_add|assert.*<', code):
                issues.append({
                    'type': 'integer_overflow',
                    'message': '📈 整数溢出风险',
                    'severity': 'high',
                    'description': 'u256加法操作可能导致整数溢出',
                    'suggestion': '使用安全的数学运算或添加溢出检查',
                    'cwe': 'CWE-190'
                })
        
        # 9. 检测函数可见性
        if re.search(r'fn\s+\w+.*\{', code) and not re.search(r'#\[external|#\[view|pub\s+fn|internal', code):
            issues.append({
                'type': 'unclear_visibility',
                'message': '👁️ 函数可见性不明确',
                'severity': 'low',
                'description': '函数可见性不明确，可能导致意外的访问',
                'suggestion': '明确指定函数可见性：#[external(v0)]、#[view] 或 internal',
                'cwe': 'CWE-668'
            })
        
        # 10. 检测gas优化问题
        if code.count('.read(') > 2:
            issues.append({
                'type': 'gas_optimization',
                'message': '⛽ Gas优化建议',
                'severity': 'info',
                'description': '多次读取存储可能消耗过多gas',
                'suggestion': '考虑将存储读取结果缓存到局部变量',
                'cwe': 'CWE-405'
            })
        
        return issues

    def assess_code_quality(self, code: str) -> Dict[str, Any]:
        """评估代码质量"""
        quality_metrics = {
            'readability_score': 0,
            'maintainability_score': 0,
            'documentation_score': 0,
            'best_practices_score': 0,
            'overall_score': 0,
            'recommendations': []
        }
        
        # 1. 可读性评估
        readability_score = 100
        lines = code.split('\n')
        
        # 检查注释覆盖率
        comment_lines = sum(1 for line in lines if line.strip().startswith('//') or line.strip().startswith('/*'))
        code_lines = sum(1 for line in lines if line.strip() and not line.strip().startswith('//'))
        comment_ratio = comment_lines / max(code_lines, 1)
        
        if comment_ratio < 0.1:
            readability_score -= 20
            quality_metrics['recommendations'].append('增加代码注释以提高可读性')
        
        # 检查变量命名
        if re.search(r'\b[a-z]\b|\b\w{1,2}\b', code):
            readability_score -= 15
            quality_metrics['recommendations'].append('使用更具描述性的变量名')
        
        # 检查函数长度
        for line in lines:
            if line.strip().startswith('fn '):
                # 简单估算函数长度
                if code.count('{') > 3:  # 嵌套过深
                    readability_score -= 10
                    quality_metrics['recommendations'].append('考虑拆分复杂函数')
                break
        
        quality_metrics['readability_score'] = max(0, readability_score)
        
        # 2. 可维护性评估
        maintainability_score = 100
        
        # 检查硬编码值
        if re.search(r'\b\d{3,}\b', code):  # 大数字硬编码
            maintainability_score -= 15
            quality_metrics['recommendations'].append('避免硬编码数值，使用常量')
        
        # 检查错误处理
        if 'fn ' in code and not re.search(r'Result|Option|assert|panic', code):
            maintainability_score -= 20
            quality_metrics['recommendations'].append('添加适当的错误处理')
        
        quality_metrics['maintainability_score'] = max(0, maintainability_score)
        
        # 3. 文档评估
        documentation_score = 100
        
        # 检查函数文档
        if 'fn ' in code and not re.search(r'///|/\*\*', code):
            documentation_score -= 30
            quality_metrics['recommendations'].append('为函数添加文档注释')
        
        # 检查参数文档
        if re.search(r'fn\s+\w+\([^)]+\)', code) and not re.search(r'@param|///.*param', code):
            documentation_score -= 20
            quality_metrics['recommendations'].append('为函数参数添加说明')
        
        quality_metrics['documentation_score'] = max(0, documentation_score)
        
        # 4. 最佳实践评估
        best_practices_score = 100
        
        # 检查是否使用了推荐的模式
        if '#[external(v0)]' in code:
            if not re.search(r'get_caller_address', code):
                best_practices_score -= 25
                quality_metrics['recommendations'].append('外部函数应验证调用者身份')
        
        # 检查是否有适当的事件
        if re.search(r'storage.*write|\.write\(', code):
            if not re.search(r'emit|event', code):
                best_practices_score -= 20
                quality_metrics['recommendations'].append('状态变更应发出事件')
        
        # 检查是否使用了安全的数学运算
        if re.search(r'[+\-*/]', code) and not re.search(r'checked_|safe_', code):
            best_practices_score -= 15
            quality_metrics['recommendations'].append('考虑使用安全的数学运算')
        
        quality_metrics['best_practices_score'] = max(0, best_practices_score)
        
        # 计算总体评分
        quality_metrics['overall_score'] = (
            quality_metrics['readability_score'] * 0.25 +
            quality_metrics['maintainability_score'] * 0.25 +
            quality_metrics['documentation_score'] * 0.25 +
            quality_metrics['best_practices_score'] * 0.25
        )
        
        return quality_metrics

    def get_security_rating(self, issues: List[Dict[str, Any]]) -> Dict[str, Any]:
        """获取安全评级"""
        critical_count = sum(1 for issue in issues if issue['severity'] == 'critical')
        high_count = sum(1 for issue in issues if issue['severity'] == 'high')
        medium_count = sum(1 for issue in issues if issue['severity'] == 'medium')
        low_count = sum(1 for issue in issues if issue['severity'] == 'low')
        
        # 计算安全评分
        security_score = 100
        security_score -= critical_count * 40
        security_score -= high_count * 25
        security_score -= medium_count * 10
        security_score -= low_count * 5
        security_score = max(0, security_score)
        
        # 确定安全等级
        if critical_count > 0:
            rating = 'F - 极度危险'
            color = '🔴'
        elif high_count > 0:
            rating = 'D - 高风险'
            color = '🟠'
        elif medium_count > 2:
            rating = 'C - 中等风险'
            color = '🟡'
        elif medium_count > 0 or low_count > 3:
            rating = 'B - 低风险'
            color = '🟢'
        else:
            rating = 'A - 安全'
            color = '✅'
        
        return {
            'score': security_score,
            'rating': rating,
            'color': color,
            'critical_issues': critical_count,
            'high_issues': high_count,
            'medium_issues': medium_count,
            'low_issues': low_count
        }

    def determine_code_type(self, code: str) -> str:
        """确定代码类型"""
        if '#[contract]' in code or '@contract' in code:
            return 'contract'
        elif 'fn ' in code and 'struct ' in code:
            return 'module'
        elif 'fn ' in code:
            return 'function'
        elif 'struct ' in code:
            return 'struct'
        elif 'trait ' in code:
            return 'trait'
        elif 'impl ' in code:
            return 'impl'
        elif 'use ' in code:
            return 'import'
        else:
            return 'snippet'


# 全局分析器实例
cairo_analyzer = CairoCodeAnalyzer()


def analyze_cairo_code(state: ChatState) -> ChatState:
    """Cairo 代码分析节点 - 增强版"""
    user_text = state["messages"][-1]["content"] if state["messages"] else ""
    
    # 提取代码块
    code_blocks = cairo_analyzer.extract_code_blocks(user_text)
    
    if not code_blocks:
        # 没有代码，设置默认值
        state["cairo_context"] = CodeAnalysisResult(
            has_code=False,
            code_blocks=[],
            code_type="none",
            syntax_errors=[],
            ast_analysis=None,
            complexity_score=0.0,
            dependencies=[],
            security_issues=[]
        )
    else:
        # 分析第一个代码块（主要代码）
        main_code = code_blocks[0]
        
        # 执行各种分析
        syntax_errors = cairo_analyzer.analyze_syntax(main_code)
        ast_analysis = cairo_analyzer.analyze_ast_structure(main_code)
        complexity = cairo_analyzer.calculate_complexity(main_code)
        dependencies = cairo_analyzer.detect_dependencies(main_code)
        security_issues = cairo_analyzer.detect_security_issues(main_code)
        code_type = cairo_analyzer.determine_code_type(main_code)
        
        # 新增：代码质量评估
        quality_metrics = cairo_analyzer.assess_code_quality(main_code)
        
        # 新增：安全评级
        security_rating = cairo_analyzer.get_security_rating(security_issues)
        
        state["cairo_context"] = CodeAnalysisResult(
            has_code=True,
            code_blocks=code_blocks,
            code_type=code_type,
            syntax_errors=syntax_errors,
            ast_analysis=ast_analysis,
            complexity_score=complexity,
            dependencies=dependencies,
            security_issues=security_issues,
            quality_metrics=quality_metrics,  # 新增
            security_rating=security_rating   # 新增
        )
    
    # 设置下一步动作
    if state["cairo_context"]["has_code"]:
        if state["cairo_context"]["syntax_errors"]:
            state["next_action"] = "error_analysis"
        else:
            state["next_action"] = "knowledge_retrieval"
    else:
        state["next_action"] = "knowledge_retrieval"
    
    return state